#include <float.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/time.h>
#include <vector>
#include <cmath>
#include <dirent.h> 
#include <iostream> 
#include <fstream>

#include "rk_defines.h"
#include "rk_debug.h"
#include "rknn_api.h"

#include "taiic_vsr.h"

using namespace std;

/*
rknntoolkit_1.5.0  output_type is int8, need use zp and  scale to trans float
*/

// 获取当前时间
static inline int64_t getCurrentTimeUs()
{
  struct timeval tv;
  gettimeofday(&tv, NULL);
  return tv.tv_sec * 1000000 + tv.tv_usec;
}

// 遍历文件夹保存文件名
void EnumerateFilesInDirectory(const std::string& path, std::vector<std::string>& files)
{
  DIR* dir;
  struct dirent* ent;
  if ((dir = opendir(path.c_str())) != NULL)
  {
    while ((ent = readdir(dir)) != NULL)
    {
        // 忽略目录和非文件条目
        if (ent->d_type == DT_REG) files.push_back(ent->d_name);
    }
    closedir(dir);
  }
  else
  {
    // 无法打开目录
    perror("");
    return;
  }
}

int main(int argc, char *argv[])
{
  // 检查参数输入
  if(argc<2)
  {
    std::cerr << "Input parameters error: ./taiic_avsr /oem/usr/model/avsr.rknn video/"; 
    return 1;
  }
  std::string video_path = argv[2]; //char *转字符串

  // 遍历输入文件夹中的文件并保存文件名
  vector<string> files;
  EnumerateFilesInDirectory(video_path, files);
  cout << "Number of files in the directory: " << files.size() << std::endl;

  // 检查输出文件是否成功打开
  ofstream outfile("result.txt");
  if (!outfile.is_open())
  {
    std::cerr << "Unable to open file result.txt";
    return 1;
  }

    // 模型推理相关初始化
    VSR_TOOLKIT_MODEL_CTX_S *vsr_ctx;
    vsr_ctx = reinterpret_cast<VSR_TOOLKIT_MODEL_CTX_S *>(malloc(sizeof(VSR_TOOLKIT_MODEL_CTX_S))); // 分配内存空间

  // 批量推理
  for (const auto& file : files)
  {
    memset(vsr_ctx, 0, sizeof(VSR_TOOLKIT_MODEL_CTX_S));
    vsr_ctx->modelPath = argv[1];
    vsr_rknn_toolkit_config_init(vsr_ctx);
    vsr_rknn_toolkit_io_init(vsr_ctx);
    
    // Get Input Data From CPU
    unsigned char *input_data = NULL;
    input_data = new unsigned char[vsr_ctx->input_attrs[0].size];
    FILE *fp = fopen((video_path + file).c_str(), "rb");
    printf("==load fiel is %s==\n", (video_path + file).c_str());
    if (fp == NULL)
    {
        perror("open failed!");
        return -1;
    }
    fread(input_data, vsr_ctx->input_attrs[0].size, 1, fp); // 读取输入数据
    fclose(fp);
    if (!input_data) return -1;

    // 模型推理
    vsr_rknn_toolkit_data_refresh(vsr_ctx, input_data);
    int64_t start_us = getCurrentTimeUs();
    int ret = rknn_run(vsr_ctx->context, NULL);
    if (ret < 0)
    {
        printf("rknn run error %d\n", ret);
        return -1;
    }
    int64_t elapse_us = getCurrentTimeUs() - start_us;
    RK_LOGD("Elapse Time = %.2fms, FPS = %.2f\n", elapse_us / 1000.f, 1000.f * 1000.f / elapse_us);
    MODEL_RESULT_S result = vsr_rknn_toolkit_result_int8_opt(vsr_ctx);
    printf("==result is %d, prob is %f==\n", result.label, result.prob);

    // Destroy rknn memory
    for (uint32_t i = 0; i < vsr_ctx->io_num.n_input; ++i)
    {
        rknn_destroy_mem(vsr_ctx->context, vsr_ctx->input_mems[i]);
    }
    for (uint32_t i = 0; i < vsr_ctx->io_num.n_output; ++i)
    {
        rknn_destroy_mem(vsr_ctx->context, vsr_ctx->output_mems[i]);
    }
    // destroy
    rknn_destroy(vsr_ctx->context);
    free(input_data);

    // 结果写入文件
    outfile << file << ": result label: " << result.label << ", prob: " << result.prob << std::endl;
  }

  // 关闭文件
  outfile.close();

  return 0;
}

